#pragma once
#include "zMatrixBase.hpp"

// for Reshape(mat,row,col)

namespace zzz{
template<typename T,class Major>
class zFunctionReshapeR : public zMatrixBaseR<T,Major>
{
public:
  explicit zFunctionReshapeR(const zMatrixBaseR<T,Major>& mat, zuint row, zuint col)
  :zMatrixBaseR<T,Major>(row,col),Mat_(mat){ZCHECK_EQ(row * col, mat.rows_ * mat.cols_);}
  const T operator()(zuint r, zuint c) const {
    CheckRange(r, c);
    return Mat_[ToIndex(r,c)];
  }
  using zMatrixBaseR<T,Major>::ToIndex;
private:
  const zMatrixBaseR<T,Major> &Mat_;
};

template<typename T,class Major>
class zFunctionReshapeW : public zMatrixBaseW<T,Major>
{
public:
  explicit zFunctionReshapeW(zMatrixBaseW<T,Major>& mat, zuint row, zuint col)
  :zMatrixBaseW<T,Major>(row,col),Mat_(mat){assert(row*col==mat.rows_*mat.cols_);}
  const T operator()(zuint r, zuint c) const {
    CheckRange(r, c);
    return Mat_[ToIndex(r,c)];
  }
  T& operator()(zuint r, zuint c) {
    CheckRange(r, c);
    return Mat_[ToIndex(r,c)];
  }
//  using zMatrixBaseW<T,Major>::operator=;
  const zMatrixBaseW<T,Major>& operator=(const zMatrixBaseR<T,Major> &other)
  {
    return zMatrixBaseW<T,Major>::operator =(other);
  }
  using zMatrixBaseR<T,Major>::ToIndex;
private:
  zMatrixBaseW<T,Major> &Mat_;
};


template<typename T,class Major>
const zFunctionReshapeR<T,Major> Reshape(const zMatrixBaseR<T,Major> &mat, zuint row, zuint col)
{
  return zFunctionReshapeR<T,Major>(mat,row,col);
}

template<typename T,class Major>
zFunctionReshapeW<T,Major> Reshape(zMatrixBaseW<T,Major> &mat, zuint row, zuint col)
{
  return zFunctionReshapeW<T,Major>(mat,row,col);
}

}
